# app/database.py
import os
from contextlib import asynccontextmanager
from typing import AsyncGenerator, Union
from dotenv import load_dotenv
from sqlalchemy import text
# 异步数据库相关
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, AsyncEngine
from sqlalchemy.ext.asyncio import async_sessionmaker
from sqlalchemy.orm import declarative_base
from sqlalchemy.pool import AsyncAdaptedQueuePool
from sqlalchemy.exc import SQLAlchemyError

# MongoDB 相关
from motor.motor_asyncio import AsyncIOMotorClient
from pymongo.errors import PyMongoError
from urllib.parse import quote_plus
import logging

logger = logging.getLogger(__name__)
# 加载环境变量
load_dotenv()

# 数据库配置
DB_USER = os.getenv("DB_USER", "postgres")
DB_PASSWORD = os.getenv("DB_PASSWORD", "postgres")
DB_HOST = os.getenv("DB_HOST", "localhost")
DB_PORT = os.getenv("DB_PORT", "5432")
DB_NAME = os.getenv("DB_NAME", "workflow_db")
DB_POOL_SIZE = int(os.getenv("DB_POOL_SIZE", 20))
DB_MAX_OVERFLOW = int(os.getenv("DB_MAX_OVERFLOW", 5))

# MongoDB 配置
MONGO_URI = os.getenv("MONGO_URI", "mongodb://localhost:27017")
MONGO_DB_NAME = os.getenv("MONGO_DB_NAME", "tws")
MONGO_POOL_SIZE = int(os.getenv("MONGO_POOL_SIZE", 10))

# 声明基类
Base = declarative_base()

class DatabaseManager:
    def __init__(self):
        self._async_engine = None
        self._async_session = None
        self._mongo_client = None
        self._mongo_db = None

    async def initialize(self):
        """初始化数据库连接"""
        # await self._init_postgres()
        await self._init_mongodb()

    async def _init_postgres(self):
        """初始化 PostgreSQL 异步连接"""
        encoded_password = quote_plus(DB_PASSWORD)
        # 构建数据库连接URL (使用 asyncpg 驱动)
        DATABASE_URL = f"postgresql+asyncpg://{DB_USER}:{encoded_password}@{DB_HOST}:{DB_PORT}/{DB_NAME}"
        # logger.info(f"Database URL: {DATABASE_URL}")
        # 创建异步数据库引擎
        self._async_engine = create_async_engine(
            DATABASE_URL,
            poolclass=AsyncAdaptedQueuePool,
            pool_size=DB_POOL_SIZE,
            max_overflow=DB_MAX_OVERFLOW,
            pool_recycle=1800,
            pool_pre_ping=True,
            echo=False,
            future=True
        )

        # 创建异步会话工厂
        self._async_session = async_sessionmaker(
            bind=self._async_engine,
            autocommit=False,
            autoflush=False,
            expire_on_commit=False,
            class_=AsyncSession
        )

    async def _init_mongodb(self):
        """初始化 MongoDB 连接"""
        self._mongo_client = AsyncIOMotorClient(
            MONGO_URI,
            maxPoolSize=MONGO_POOL_SIZE,
            minPoolSize=1,
            connectTimeoutMS=5000,
            serverSelectionTimeoutMS=5000
        )

        # 测试连接
        try:
            # 发送 ping 命令检查连接
            await self._mongo_client.admin.command('ping')
            logger.info("✅ MongoDB connection successful")
            self._mongo_db = self._mongo_client[MONGO_DB_NAME]
        except PyMongoError as e:
            logger.info(f"❌ MongoDB connection failed: {str(e)}")
            self._mongo_client = None
            self._mongo_db = None

    @property
    def mongo_db(self):
        """获取 MongoDB 数据库实例"""
        if not self._mongo_db:
            raise RuntimeError("MongoDB is not initialized")
        return self._mongo_db

    @asynccontextmanager
    async def session(self) -> AsyncGenerator[AsyncSession, None]:
        """异步数据库会话上下文管理器"""
        if not self._async_session:
            raise RuntimeError("Database session is not initialized")

        async_session = self._async_session()
        try:
            yield async_session
            await async_session.commit()
        except SQLAlchemyError as exc:
            await async_session.rollback()
            raise exc
        finally:
            await async_session.close()

    async def get_db(self) -> AsyncGenerator[AsyncSession, None]:
        """用于 FastAPI 依赖注入的异步会话生成器"""
        if not self._async_session:
            raise RuntimeError("Database session is not initialized")

        async with self._async_session() as session:
            try:
                yield session
            except SQLAlchemyError as exc:
                await session.rollback()
                raise exc
            finally:
                await session.close()

    async def close(self):
        """关闭所有数据库连接"""
        if self._async_engine:
            await self._async_engine.dispose()
            logger.info("PostgreSQL connections closed")

        if self._mongo_client:
            self._mongo_client.close()
            logger.info("MongoDB connections closed")

    async def test_postgres_connection(self) -> bool:
        """测试 PostgreSQL 连接是否正常"""
        if not self._async_engine:
            return False

        try:
            async with self._async_engine.connect() as conn:
                result = await conn.execute(text("SELECT 1"))
                return result.scalar() == 1
        except SQLAlchemyError as e:
            logger.info(f"PostgreSQL connection failed: {str(e)}")
            return False

    async def test_mongodb_connection(self) -> bool:
        """测试 MongoDB 连接是否正常"""
        if not self._mongo_client:
            return False

        try:
            await self._mongo_client.admin.command('ping')
            return True
        except PyMongoError as e:
            logger.info(f"MongoDB connection failed: {str(e)}")
            return False

    async def init_db(self):
        """初始化数据库，创建所有表（仅用于开发环境）"""
        if not self._async_engine:
            raise RuntimeError("Database engine is not initialized")

        import app.models  # 导入所有模型以注册到Base.metadata

        logger.info("Initializing database...")
        async with self._async_engine.begin() as conn:
            # 创建所有表
            await conn.run_sync(Base.metadata.create_all)
        logger.info("Database initialized successfully")


# 创建全局数据库管理器实例
db_manager = DatabaseManager()


# 应用启动时初始化数据库连接
async def init_database():
    """初始化数据库连接（在应用启动时调用）"""
    await db_manager.initialize()

    # 测试连接
    postgres_ok = await db_manager.test_postgres_connection()
    mongo_ok = await db_manager.test_mongodb_connection()

    logger.info(f"{'✅' if postgres_ok else '❌'}:PostgreSQL connection")
    logger.info(f"{'✅' if mongo_ok else '❌'}:MongoDB connection")

    # 在开发环境中创建表
    if os.getenv("ENV") == "development" and postgres_ok:
        await db_manager.init_db()


# 应用关闭时清理数据库连接
async def close_database():
    """关闭数据库连接（在应用关闭时调用）"""
    await db_manager.close()


# 获取 MongoDB 集合的辅助函数
def get_mongo_collection(collection_name: str):
    """获取 MongoDB 集合实例"""
    return db_manager.mongo_db[collection_name]


# 获取 PostgreSQL 会话的依赖函数
async def get_db() -> AsyncGenerator[AsyncSession, None]:
    """用于 FastAPI 依赖注入的数据库会话"""
    async for session in db_manager.get_db():
        yield session